'''
The \ALT command
----------------

The \ALT command is used to define different alternatives in the source 
file. The following code generates a template that offers three different 
alternatives "a", "b", or "c"

>>> alt = AltSource(r'\ALT{a|b|c}')

These alternatives can be accessed by index

>>> alt[0], alt[1], alt[2] 
('a', 'b', 'c')

All AltSource objects have at least one section called 'default'. It is 
possible to define alternatives in different sections by using an optional
argument in the template.

>>> alt = AltSource(r'\ALT[foo, bar]{a|b|c}')

Documents in each section can be accessed by a combination of <section name>
and <index>

>>> alt['foo', 0]
'a'
>>> alt['default', 0]
''

If only the section name is used it returns a tuple of all alternatives for
that section

>>> alt['bar']
('a', 'b', 'c')

If no section is specified in an \ALT command, the alternatives are 
replicated in all sections

>>> alt = AltSource(r'\ALT[bar]{a|b}\ALT{0|1}')
>>> alt['bar', 0], alt[0] # alt[0] <==> alt['default', 0]
('a0', '0')

Whitespace is removed from alternatives. This can be convenient to format 
alternatives in more readable ways

>>> src = r"""\ALT{ 
...     foo
...     |bar
...     |foobar
... }"""
>>> AltSource(src)['default']
('foo', 'bar', 'foobar')
 
'''
import re
import itertools
from collections import Mapping
from alttex.util import lcm, readfile
from alttex.texutil import LaTeXError, read_group, read_altarg, read_any

class AltSource(Mapping):
    r'''An AltSource object works like a dictionary of "(section, idx): document" 
    pairs generated by the given template 'source'.
    
    Parameters
    ==========
    
    source : str or file
        A string holding the template source.
    filename : str
        An optional filename. If not given and source is a file, it is 
        automatically extracted from the ".name" attribute.
    '''
    _cmd_regex = re.compile(r'(?P<name>\\ALT|\\IF|\\begin)')

    def __init__(self, source, filename=None):
        self.source, self.filename = readfile(source, None)
        if filename is not None:
            self.filename = filename

        self._cache = {}                    # cache of rendered documents
        self._sections = set(['default'])   # store section names
        self._sizes = set([1])              # store size of \ALT blocks

        # Parse document and save parsed items
        self._parsed = list(self._fix_else(self._parse()))

    #===========================================================================
    # Parsing functions
    #===========================================================================
    def _parse(self):
        '''Yields strings for LaTeX sections and dictionaries mapping sections
        to a list of alternatives.'''

        S = self.source
        pos = 0
        while True:
            # Find next command
            pre, cmd, pos = read_any(S, (r'\ALT', r'\IF', r'\ELSE'), pos)
            yield pre
            if cmd is None:
                return

            # Process \ALT command ---------------------------------------------
            if cmd == '\\ALT':
                # Read sections
                sections, pos = read_altarg(S, pos, retpos=True, strip=True)
                if sections is not None:
                    sections = [sec.strip() for sec in sections.split(',')]
                else:
                    sections = [None]
                self._sections.update(sections)

                # Read body
                body, pos = read_group(S, pos, retpos=True, strip=True)
                body = tuple(b.strip() for b in body.split('|'))
                self._sizes.add(len(body))
                out = { sec: body for sec in (sections or ['default']) }
                out['@cmd'] = 'ALT'
                yield out

            # Process \IF command ----------------------------------------------
            elif cmd == '\\IF':
                # Read sections
                sections, pos = read_group(S, pos, retpos=True, strip=True)
                sections = [sec.strip() for sec in sections.split(',')]
                self._sections.update(sections)

                # Read body
                body, pos = read_group(S, pos, retpos=True, strip=True)
                if r'\IF' in body:
                    raise LaTeXError(r'\IF cannot be inside another \IF block', S, pos)
                if r'\ELSE' in body:
                    raise LaTeXError(r'\ELSE cannot be inside an \IF block', S, pos)
                if r'\ALT' in body:
                    alt = AltSource(body)
                    self._sizes.update(alt._sizes)
                    docs = alt.default
                    out = { sec: docs for sec in sections }
                elif sections != [None]:
                    out = { sec: (body,) for sec in sections }
                else:
                    continue
                out['@cmd'] = 'IF'
                yield out

            # Process \ELSE command ----------------------------------------------
            elif cmd == '\\ELSE':
                body, pos = read_group(S, pos, retpos=True, strip=True)
                if r'\ALT' in body:
                    docs = self._alt_sub_block(body, pos)
                    out = {'@ELSE': docs}
                else:
                    out = {'@ELSE': (body,)}
                out['@cmd'] = 'ELSE'
                out['@pos'] = pos
                yield out

    def _alt_sub_block(self, body, pos=0):
        alt = AltSource(body)
        self._sizes.update(alt._sizes)
        return alt.default

    def _fix_else(self, it):
        'Executed after parsing to expand the ELSE blocks'

        L = list(it)
        sections = set(self.sections())
        for i, obj in enumerate(L):
            if isinstance(obj, dict) and obj['@cmd'] == 'ELSE':
                prev = L[i - 1]
                if isinstance(prev, str):
                    prev = L[i - 2]
                if not prev['@cmd'] == 'IF':
                    raise LaTeXError(r'\ELSE must appear after an \IF command', self.source, obj['@pos'])
                for sec in prev:
                    sections.discard(sec)
                docs = obj['@ELSE']
                obj.update({sec: docs for sec in sections})
            yield obj

    #===========================================================================
    # Magic functions
    #===========================================================================
    def __getitem__(self, item):
        try:
            return self._cache[item]
        except KeyError:
            # Retrieve list of alternatives by name
            if isinstance(item, str):
                self._cache[item] = self.get_section(item)

            # Retrieve 'default' by number
            elif isinstance(item, int):
                self._cache[item] = self['default', item]

            # Retrieve by name, number
            else:
                name, idx = item
                max_size = lcm(self._sizes)
                if idx < max_size:
                    self._cache[item] = self[name][idx]
                else:
                    return self[name, idx - max_size]

            return self._cache[item]

    def __iter__(self):
        for sec in self._sections:
            if sec is not None:
                yield sec

    def __len__(self):
        return len(self._sections)

    #===========================================================================
    # API
    #===========================================================================
    def sections(self):
        '''Return a list of sections in the document'''

        return list(iter(self))

    def _fill_list(self, L, size):
        '''Return a list of given size cycling over the elements of L. Used by 
        the get_section() method'''

        seq = itertools.cycle(L)
        seq = itertools.islice(seq, size)
        return list(seq)

    def get_section(self, section, size=None):
        '''Returns a list of `size` documents in the given `section`.'''

        if size is None:
            size = lcm(self._sizes)
        elems = []
        for obj in self._parsed:
            if isinstance(obj, str):
                L = [obj] * size
            else:
                try:
                    L = obj[section]
                except KeyError:
                    L = obj.get(None, [''])
                L = self._fill_list(L, size)
            elems.append(L)

        return tuple(''.join(obj[i] for obj in elems) for i in range(size))

    @property
    def default(self):
        return self['default']

if __name__ == '__main__':
    import doctest
    doctest.testmod()
